package com.cfp4cloud.cfp.knowledge.support.handler.parse;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.lang.Pair;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.cfp4cloud.cfp.admin.api.utils.ByteArrayMultipartFile;
import com.cfp4cloud.cfp.common.core.constant.enums.YesNoEnum;
import com.cfp4cloud.cfp.knowledge.config.properties.AiKnowledgeProperties;
import com.cfp4cloud.cfp.knowledge.dto.AiMarkitdownDTO;
import com.cfp4cloud.cfp.knowledge.dto.MarkitdownResponseDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiDatasetEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiDocumentEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiModelEntity;
import com.cfp4cloud.cfp.knowledge.mapper.AiModelMapper;
import com.cfp4cloud.cfp.knowledge.service.AiMarkitdownAssistantService;
import com.cfp4cloud.cfp.knowledge.support.constant.FileParserStatusEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.FileTypeEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.ModelTypeEnums;
import com.cfp4cloud.cfp.knowledge.support.util.PromptBuilder;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.InputStream;
import java.util.List;
import java.util.Objects;

/**
 * Office文档转Markdown解析处理器
 * <p>
 * 该处理器使用Markitdown服务将Office文档（DOCX、PDF、PPTX、XLSX）转换为Markdown格式。
 * Markitdown服务可以保留文档的结构和格式信息，转换后的Markdown文本更适合进行后续的文本分析和处理。 如果配置了视觉模型，还可以对文档中的图片进行OCR识别。
 *
 * @author chenda
 * @date 2024/12/22
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class Office2MdUploadFileParseHandler implements UploadFileParseHandler {

	private final AiKnowledgeProperties aiKnowledgeProperties;

	private final AiMarkitdownAssistantService markitdownAssistantService;

	private final AiModelMapper aiModelMapper;

	/**
	 * 判断是否支持处理指定类型的文件
	 * <p>
	 * 首先检查Markitdown服务是否启用，如果未启用则不支持任何文件。
	 * 当Markitdown服务启用时，支持处理DOCX、PDF、PPTX、XLSX格式的Office文档。
	 * @param aiDataset 数据集实体
	 * @param documentEntity 文档实体，包含文件信息
	 * @return 如果支持该文件类型返回true，否则返回false
	 */
	@Override
	public boolean supports(AiDatasetEntity aiDataset, AiDocumentEntity documentEntity) {
		String contentType = FileUtil.extName(documentEntity.getFileUrl());
		return aiKnowledgeProperties.getMarkitdown().isEnabled() && (FileTypeEnums.DOCX.getType().equals(contentType)
				|| FileTypeEnums.PDF.getType().equals(contentType) || FileTypeEnums.PPTX.getType().equals(contentType)
				|| FileTypeEnums.XLSX.getType().equals(contentType));
	}

	/**
	 * 将Office文档转换为Markdown格式文本
	 * <p>
	 * 该方法执行以下步骤： 1. 将输入流转换为ByteArrayMultipartFile对象 2. 查询默认的视觉模型（如果有），用于处理文档中的图片 3.
	 * 如果找到视觉模型，设置模型参数和OCR提示词 4. 调用Markitdown服务进行文档转换
	 * @param documentEntity 文档实体，包含文档信息
	 * @param inputStream 文件输入流
	 * @param extName 文件扩展名
	 * @return 包含解析状态和Markdown文本的Pair对象，成功时返回转换后的文本，失败时返回错误信息
	 */
	@Override
	public Pair<FileParserStatusEnums, String> file2String(AiDocumentEntity documentEntity, InputStream inputStream,
			String extName) {
		try {

			ByteArrayMultipartFile file = new ByteArrayMultipartFile(documentEntity.getName(), documentEntity.getName(),
					null, inputStream.readAllBytes());

			AiMarkitdownDTO request = new AiMarkitdownDTO();

			// 查询默认的视觉模型
			AiModelEntity aiModelEntity = aiModelMapper.selectOne(Wrappers.<AiModelEntity>lambdaQuery()
				.eq(AiModelEntity::getModelType, ModelTypeEnums.VISION.getType())
				.eq(AiModelEntity::getDefaultModel, YesNoEnum.YES.getCode()), false);

			if (Objects.nonNull(aiModelEntity)) {
				request.setModel(aiModelEntity.getModelName());
				request.setApi_key(aiModelEntity.getApiKey());
				request.setBase_url(aiModelEntity.getBaseUrl());
				request.setPrompt(PromptBuilder.render("ocr-image.st"));
			}

			MarkitdownResponseDTO markitdownResponseDTO = markitdownAssistantService.upload(file, request);
			return Pair.of(FileParserStatusEnums.PARSE_SUCCESS, markitdownResponseDTO.getText());
		}
		catch (Exception e) {
			log.error("文件 {} 解析失败", documentEntity.getName(), e);
			return Pair.of(FileParserStatusEnums.PARSE_FAIL, e.getMessage());
		}
	}

	/**
	 * 获取处理器的优先级顺序
	 * <p>
	 * 数值越大，优先级越高。该处理器的优先级为100， 在多个处理器都支持同一文件类型时的加载顺序。
	 * @return 优先级值，返回100
	 */
	@Override
	public int getOrder() {
		return 100;
	}

}
